import torch
import torch.nn as nn
import torch.optim as optim
import yaml
from tqdm import tqdm

from data.loader import prepare_data
from models.resnet_se import ResNetSE


def load_config():
    with open("C:/Users/23639/PycharmProjects/TCM/TCM/configs/default.yaml") as f:
        return yaml.safe_load(f)


def train_model():
    config = load_config()
    device = torch.device(config['training']['device'] if torch.cuda.is_available() else "cpu")

    # 加载数据
    train_loader, val_loader = prepare_data(config)

    model = ResNetSE(
        num_classes=config['model']['num_classes'],
        pretrained=config['model']['pretrained']
    )
    model = model.to(device)

    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config['training']['lr'])

    # 训练过程
    for epoch in range(config['training']['epochs']):
        # 训练
        model.train()
        running_loss, running_correct = 0.0, 0
        for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}"):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            running_correct += torch.sum(predicted == labels.data)

        # 验证
        model.eval()
        val_loss, val_correct = validate(model, val_loader, criterion, device)

        # 打印
        print_metrics(train_loader, val_loader, epoch, running_loss, running_correct, val_loss, val_correct,
                      len(train_loader), len(val_loader.dataset))

    # 保存结果
    torch.save(model.state_dict(), "tcm_resnet_se.pth")


def validate(model, val_loader, criterion, device):
    model.eval()
    val_loss, val_correct = 0.0, 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            val_loss += criterion(outputs, labels).item()
            _, predicted = torch.max(outputs, 1)
            val_correct += torch.sum(predicted == labels.data)
    return val_loss, val_correct


def print_metrics(train_loader, val_loader, epoch, train_loss, train_correct, val_loss, val_correct, num_batches,
                  val_size):
    # 转换为浮点型
    train_correct = train_correct.float() if torch.is_tensor(train_correct) else float(train_correct)
    val_correct = val_correct.float() if torch.is_tensor(val_correct) else float(val_correct)

    train_loss /= num_batches
    train_acc = train_correct / (num_batches * train_loader.batch_size)

    val_loss /= len(val_loader)
    val_acc = val_correct / val_size

    print(f"Epoch {epoch + 1} | "
          f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
          f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")


if __name__ == "__main__":
    train_model()
